package brainfsck

import chisel3._
import chisel3.util._

object BrainfsckInstr extends ChiselEnum {
    val SHIFT_RIGHT         = Value(0.U(3.W))
    val SHIFT_LEFT          = Value(1.U(3.W))
    val INCREASE            = Value(2.U(3.W))
    val DECREASE            = Value(3.U(3.W))
    val OUTPUT              = Value(4.U(3.W))
    val INPUT               = Value(5.U(3.W))
    val BRANCH_RIGHT_ZERO   = Value(6.U(3.W))
    val BRANCH_LEFT_NONZERO = Value(7.U(3.W))
}

class DecodedBrainfsckInstr extends Bundle {
    val branchLeftNonZero = Bool()
    val branchRightZero   = Bool()
    val input  = Bool()
    val output = Bool()
    val decrease = Bool()
    val increase = Bool()
    val shiftLeft  = Bool()
    val shiftRight = Bool()
}

class Brainfsck(pcWidth: Int) extends Module {
    val io = IO(new Bundle {
        val memPort = Flipped(new DataMemPort)
        val inputPort = Flipped(Decoupled(UInt(8.W)))
        val outputPort = Decoupled(UInt(8.W))
        val romPort = Flipped(new RomPort(pcWidth))
    })

    private val decodedInstr = UIntToOH(io.romPort.data.asUInt).asTypeOf(new DecodedBrainfsckInstr)

    private val loopLevel = RegInit(0.S(5.W))
    private val instrEnable = loopLevel === 0.S
    loopLevel := MuxCase(loopLevel, Seq(
        (decodedInstr.branchRightZero && (!instrEnable || io.memPort.readValue===0.U)) -> (loopLevel + 1.S),
        (decodedInstr.branchLeftNonZero && (!instrEnable || io.memPort.readValue=/=0.U)) -> (loopLevel - 1.S),
    ))

    private val advance = !instrEnable || Mux1H(Seq(
        decodedInstr.shiftLeft -> true.B,
        decodedInstr.shiftRight -> true.B,
        decodedInstr.increase -> true.B,
        decodedInstr.decrease -> true.B,
        decodedInstr.output -> io.outputPort.ready,
        decodedInstr.input -> io.inputPort.valid,
        decodedInstr.branchRightZero -> true.B,
        decodedInstr.branchLeftNonZero -> true.B,
    ))

    io.memPort.shiftLeft := instrEnable && decodedInstr.shiftLeft
    io.memPort.shiftRight := instrEnable && decodedInstr.shiftRight

    io.memPort.writeEnable := instrEnable && advance && (decodedInstr.increase || decodedInstr.decrease || decodedInstr.input)
    io.memPort.writeValue := Mux1H(Seq(
        (decodedInstr.input && instrEnable && advance) -> io.inputPort.bits,
        (decodedInstr.increase && instrEnable && advance) -> (io.memPort.readValue + 1.U),
        (decodedInstr.decrease && instrEnable && advance) -> (io.memPort.readValue - 1.U),
    ))

    io.memPort.shiftLeft := instrEnable && decodedInstr.shiftLeft
    io.memPort.shiftRight := instrEnable && decodedInstr.shiftRight

    private val pc = RegInit(0.U(pcWidth.W))
    private val nextPc = MuxCase(pc + 1.U, Seq(
        (loopLevel===(-1).S && decodedInstr.branchRightZero) -> (pc + 1.U),
        (loopLevel===0.S && decodedInstr.branchLeftNonZero && io.memPort.readValue=/=0.U) -> (pc - 1.U),
        (loopLevel < 0.S) -> (pc - 1.U),
    ))
    pc := Mux(advance, nextPc, pc)
    io.romPort.ce := Mux(reset.asBool, true.B, advance)
    io.romPort.addr := Mux(reset.asBool, 0.U(pcWidth.W), nextPc)

    io.outputPort.bits := io.memPort.readValue
    io.outputPort.valid := instrEnable && decodedInstr.output

    io.inputPort.ready := instrEnable && decodedInstr.input
}
